"""
Game state management system for tank battle game
"""
import pygame
import sys
import os
from typing import Dict, List, Optional, Callable

# Add src directory to path for imports
current_dir = os.path.dirname(os.path.abspath(__file__))
src_dir = os.path.dirname(current_dir)
if src_dir not in sys.path:
    sys.path.insert(0, src_dir)

from utils.constants import GameState


class GameStateManager:
    """Game state management system"""
    
    def __init__(self):
        """Initialize game state manager"""
        self.current_state = GameState.MENU
        self.previous_state = None
        self.state_transitions = {}
        self.state_handlers = {}
        self.state_data = {}
        
        # Initialize state handlers
        self._initialize_state_handlers()
    
    def _initialize_state_handlers(self):
        """Initialize state transition handlers"""
        # Menu state handlers
        self.state_handlers[GameState.MENU] = {
            'enter': self._on_enter_menu,
            'exit': self._on_exit_menu,
            'update': self._on_update_menu,
            'render': self._on_render_menu
        }
        
        # Playing state handlers
        self.state_handlers[GameState.PLAYING] = {
            'enter': self._on_enter_playing,
            'exit': self._on_exit_playing,
            'update': self._on_update_playing,
            'render': self._on_render_playing
        }
        
        # Paused state handlers
        self.state_handlers[GameState.PAUSED] = {
            'enter': self._on_enter_paused,
            'exit': self._on_exit_paused,
            'update': self._on_update_paused,
            'render': self._on_render_paused
        }
        
        # Game over state handlers
        self.state_handlers[GameState.GAME_OVER] = {
            'enter': self._on_enter_game_over,
            'exit': self._on_exit_game_over,
            'update': self._on_update_game_over,
            'render': self._on_render_game_over
        }
        
        # Settings state handlers
        self.state_handlers[GameState.SETTINGS] = {
            'enter': self._on_enter_settings,
            'exit': self._on_exit_settings,
            'update': self._on_update_settings,
            'render': self._on_render_settings
        }
    
    def change_state(self, new_state: str, data: Optional[Dict] = None):
        """
        Change game state
        
        Args:
            new_state: New state to transition to
            data: Optional data to pass to new state
        """
        if new_state == self.current_state:
            return
        
        # Validate state transition
        if not self._is_valid_transition(self.current_state, new_state):
            return
        
        # Exit current state
        if self.current_state in self.state_handlers:
            self.state_handlers[self.current_state]['exit']()
        
        # Store previous state
        self.previous_state = self.current_state
        
        # Change state
        self.current_state = new_state
        
        # Store state data
        if data:
            self.state_data[new_state] = data
        
        # Enter new state
        if new_state in self.state_handlers:
            self.state_handlers[new_state]['enter']()
    
    def _is_valid_transition(self, from_state: str, to_state: str) -> bool:
        """Check if state transition is valid"""
        # Define valid transitions
        valid_transitions = {
            GameState.MENU: [GameState.PLAYING, GameState.SETTINGS],
            GameState.PLAYING: [GameState.PAUSED, GameState.GAME_OVER],
            GameState.PAUSED: [GameState.PLAYING, GameState.MENU],
            GameState.GAME_OVER: [GameState.MENU, GameState.PLAYING],
            GameState.SETTINGS: [GameState.MENU]
        }
        
        return to_state in valid_transitions.get(from_state, [])
    
    def get_current_state(self) -> str:
        """Get current state"""
        return self.current_state
    
    def get_previous_state(self) -> Optional[str]:
        """Get previous state"""
        return self.previous_state
    
    def update(self, dt: float):
        """Update current state"""
        if self.current_state in self.state_handlers:
            self.state_handlers[self.current_state]['update'](dt)
    
    def render(self, screen: pygame.Surface):
        """Render current state"""
        if self.current_state in self.state_handlers:
            self.state_handlers[self.current_state]['render'](screen)
    
    def handle_event(self, event: pygame.event.Event) -> bool:
        """
        Handle pygame event
        
        Args:
            event: Pygame event
            
        Returns:
            True if event was handled
        """
        # Common event handling
        if event.type == pygame.KEYDOWN:
            if event.key == pygame.K_ESCAPE:
                return self._handle_escape_key()
        
        # State-specific event handling
        if self.current_state in self.state_handlers:
            handler = self.state_handlers[self.current_state].get('handle_event')
            if handler:
                return handler(event)
        
        return False
    
    def _handle_escape_key(self) -> bool:
        """Handle escape key press"""
        if self.current_state == GameState.PLAYING:
            self.change_state(GameState.PAUSED)
            return True
        elif self.current_state == GameState.PAUSED:
            self.change_state(GameState.PLAYING)
            return True
        elif self.current_state == GameState.SETTINGS:
            self.change_state(GameState.MENU)
            return True
        
        return False
    
    # State handler methods
    def _on_enter_menu(self):
        """Called when entering menu state"""
        pass
    
    def _on_exit_menu(self):
        """Called when exiting menu state"""
        pass
    
    def _on_update_menu(self, dt: float):
        """Update menu state"""
        pass
    
    def _on_render_menu(self, screen: pygame.Surface):
        """Render menu state"""
        pass
    
    def _on_enter_playing(self):
        """Called when entering playing state"""
        pass
    
    def _on_exit_playing(self):
        """Called when exiting playing state"""
        pass
    
    def _on_update_playing(self, dt: float):
        """Update playing state"""
        pass
    
    def _on_render_playing(self, screen: pygame.Surface):
        """Render playing state"""
        pass
    
    def _on_enter_paused(self):
        """Called when entering paused state"""
        pass
    
    def _on_exit_paused(self):
        """Called when exiting paused state"""
        pass
    
    def _on_update_paused(self, dt: float):
        """Update paused state"""
        pass
    
    def _on_render_paused(self, screen: pygame.Surface):
        """Render paused state"""
        pass
    
    def _on_enter_game_over(self):
        """Called when entering game over state"""
        pass
    
    def _on_exit_game_over(self):
        """Called when exiting game over state"""
        pass
    
    def _on_update_game_over(self, dt: float):
        """Update game over state"""
        pass
    
    def _on_render_game_over(self, screen: pygame.Surface):
        """Render game over state"""
        pass
    
    def _on_enter_settings(self):
        """Called when entering settings state"""
        pass
    
    def _on_exit_settings(self):
        """Called when exiting settings state"""
        pass
    
    def _on_update_settings(self, dt: float):
        """Update settings state"""
        pass
    
    def _on_render_settings(self, screen: pygame.Surface):
        """Render settings state"""
        pass
    
    def set_state_handler(self, state: str, handler_type: str, handler: Callable):
        """
        Set custom state handler
        
        Args:
            state: State name
            handler_type: Type of handler ('enter', 'exit', 'update', 'render', 'handle_event')
            handler: Handler function
        """
        if state not in self.state_handlers:
            self.state_handlers[state] = {}
        
        self.state_handlers[state][handler_type] = handler
    
    def get_state_data(self, state: Optional[str] = None) -> Optional[Dict]:
        """Get data for specified state or current state"""
        if state is None:
            state = self.current_state
        
        return self.state_data.get(state)
    
    def set_state_data(self, state: str, data: Dict):
        """Set data for specified state"""
        self.state_data[state] = data
    
    def clear_state_data(self, state: Optional[str] = None):
        """Clear data for specified state or all states"""
        if state:
            self.state_data.pop(state, None)
        else:
            self.state_data.clear()
